#include "VStdAfx.h"
#include "VTextureAtlas.h"

VTextureAtlas::VTextureAtlas(UINT InitWidth, UINT InitHeight, UINT MaxWidth, UINT MaxHeight, BOOL PowerOfTwo):
m_Width(InitWidth),
	m_Height(InitHeight),
	m_PowOfTwo(PowerOfTwo)
{
	// push root node
	m_Nodes.PushBack(AtlasNode(0,0,MaxWidth,MaxHeight));
}


VTextureAtlas::~VTextureAtlas(void)
{

}

void VTextureAtlas::Clear(UINT MaxWidth, UINT MaxHeight)
{
	m_Nodes.Clear();
	m_Nodes.PushBack(AtlasNode(0,0,MaxWidth,MaxHeight));
}

/** Recursively traverses the tree depth first and searches for a large enough leaf node to contain the requested allocation. */
int VTextureAtlas::TryAddChildNode(int NodeIndex, UINT Width, UINT Height, BOOL bAllowTextureEnlargement)
{
	// Store a copy of the current node on the stack for easier debugging.
	// Can't store a pointer to the current node since Nodes may be reallocated in this function.
	const AtlasNode CurrentNode = m_Nodes[NodeIndex];
	// But do access this node via a pointer until the first recursive call. Prevents a ton of LHS.
	const AtlasNode* CurrentNodePtr = &m_Nodes[NodeIndex];
	if (CurrentNodePtr->ChildA != -1)
	{
		// Children are always allocated together
		VASSERT(CurrentNodePtr->ChildB != -1);
		// Traverse the children
		const int Result = TryAddChildNode(CurrentNodePtr->ChildA, Width, Height, bAllowTextureEnlargement);
		// The pointer is now invalid, be explicit!
		CurrentNodePtr = 0;

		if (Result != -1)
		{
			return Result;
		}

		return TryAddChildNode(CurrentNode.ChildB, Width, Height, bAllowTextureEnlargement);
	}
	// AtlasNode has no children, it is a leaf
	else
	{
		// Reject this node if it is already used
		if (CurrentNodePtr->bUsed)
		{
			return -1;
		}

		// Reject this node if it is too small for the element being placed
		if (CurrentNodePtr->Width < Width || CurrentNodePtr->Height < Height)
		{
			return -1;
		}

		if (!bAllowTextureEnlargement)
		{
			// Reject this node if this is an attempt to allocate space without enlarging the texture, 
			// And this node cannot hold the element without enlarging the texture.
			if (CurrentNodePtr->OffsetX + Width > m_Width || CurrentNodePtr->OffsetY + Height > m_Height)
			{
				return -1;
			}
		}

		// Use this node if the size matches the requested element size
		if (CurrentNodePtr->Width == Width && CurrentNodePtr->Height == Height)
		{
			return NodeIndex;
		}

		const UINT ExcessWidth = CurrentNodePtr->Width - Width;
		const UINT ExcessHeight = CurrentNodePtr->Height - Height;

		// The pointer to the current node may be invalidated below, be explicit!
		CurrentNodePtr = 0;

		// Add new nodes, and link them as children of the current node.
		if (ExcessWidth > ExcessHeight)
		{
			// Update the child indices
			m_Nodes[NodeIndex].ChildA = m_Nodes.Size();

			// Create a child with the same width as the element being placed.
			// The height may not be the same as the element height yet, in that case another subdivision will occur when traversing this child node.
			m_Nodes.PushBack(AtlasNode(
				CurrentNode.OffsetX,
				CurrentNode.OffsetY,
				Width,
				CurrentNode.Height
				));

			// Create a second child to contain the leftover area in the X direction
			m_Nodes[NodeIndex].ChildB = m_Nodes.Size();
			m_Nodes.PushBack(AtlasNode(
				CurrentNode.OffsetX + Width,
				CurrentNode.OffsetY,
				CurrentNode.Width - Width,
				CurrentNode.Height
				));
		}
		else
		{
			m_Nodes[NodeIndex].ChildA = m_Nodes.Size();
			m_Nodes.PushBack(AtlasNode(
				CurrentNode.OffsetX,
				CurrentNode.OffsetY,
				CurrentNode.Width,
				Height
				));

			m_Nodes[NodeIndex].ChildB = m_Nodes.Size();
			m_Nodes.PushBack(AtlasNode(
				CurrentNode.OffsetX,
				CurrentNode.OffsetY + Height,
				CurrentNode.Width,
				CurrentNode.Height - Height
				));
		}

		// Only traversing ChildA, since ChildA is always the newly created node that matches the element size
		return TryAddChildNode(m_Nodes[NodeIndex].ChildA, Width, Height, bAllowTextureEnlargement);
	}
}

BOOL VTextureAtlas::AddElement(UINT& OffsetX, UINT& OffsetY, UINT Width, UINT Height)
{
	if (Width == 0 || Height == 0)
	{
		OffsetX = 0;
		OffsetY = 0;
		return TRUE;
	}

	// Try allocating space without enlarging the texture.
	int	NodeIndex = TryAddChildNode(0, Width, Height, FALSE);
	if (NodeIndex == -1)
	{
		// Try allocating space which might enlarge the texture.
		NodeIndex = TryAddChildNode(0, Width, Height, TRUE);
	}

	if (NodeIndex != -1)
	{
		AtlasNode&	Node = m_Nodes[NodeIndex];
		Node.bUsed = TRUE;
		OffsetX = Node.OffsetX;
		OffsetY = Node.OffsetY;

		if (m_PowOfTwo)
		{
			m_Width = VMax<UINT>(m_Width, VNextPOT(Node.OffsetX + Width));
			m_Height = VMax<UINT>(m_Height, VNextPOT(Node.OffsetX + Height));
		}
		else
		{
			m_Width = VMax<UINT>(m_Width, Node.OffsetX + Width);
			m_Height = VMax<UINT>(m_Height, Node.OffsetX + Height);
		}
		return TRUE;
	}
	else
	{
		return FALSE;
	}
}

/** Recursively removes the children of a given node from the Nodes array and adjusts existing indices to compensate. */
void VTextureAtlas::RemoveChildNode(INT NodeIndex)
{
	// Traverse the children depth first
	if (m_Nodes[NodeIndex].ChildA != -1)
	{
		RemoveChildNode(m_Nodes[NodeIndex].ChildA);
	}
	if (m_Nodes[NodeIndex].ChildB != -1)
	{
		RemoveChildNode(m_Nodes[NodeIndex].ChildB);
	}

	if (m_Nodes[NodeIndex].ChildA != -1)
	{
		// Store off the index of the child since it may be changed in the code below
		const INT OldChildA = m_Nodes[NodeIndex].ChildA;

		// Remove the child from the Nodes array
		m_Nodes.EraseAt(OldChildA);

		// Iterate over all the Nodes and fix up their indices now that an element has been removed
		for (int OtherNodeIndex = 0; OtherNodeIndex < (int)m_Nodes.Size(); OtherNodeIndex++)
		{
			if (m_Nodes[OtherNodeIndex].ChildA >= OldChildA)
			{
				m_Nodes[OtherNodeIndex].ChildA--;
			}
			if (m_Nodes[OtherNodeIndex].ChildB >= OldChildA)
			{
				m_Nodes[OtherNodeIndex].ChildB--;
			}
		}
		// Mark the node as not having a ChildA
		m_Nodes[NodeIndex].ChildA = -1;
	}

	if (m_Nodes[NodeIndex].ChildB != -1)
	{
		const int OldChildB = m_Nodes[NodeIndex].ChildB;
		m_Nodes.EraseAt(OldChildB);
		for (INT OtherNodeIndex = 0; OtherNodeIndex < (int)m_Nodes.Size(); OtherNodeIndex++)
		{
			if (m_Nodes[OtherNodeIndex].ChildA >= OldChildB)
			{
				m_Nodes[OtherNodeIndex].ChildA--;
			}
			if (m_Nodes[OtherNodeIndex].ChildB >= OldChildB)
			{
				m_Nodes[OtherNodeIndex].ChildB--;
			}
		}
		m_Nodes[NodeIndex].ChildB = -1;
	}
}

BOOL VTextureAtlas::RemoveElement(UINT OffsetX, UINT OffsetY, UINT Width, UINT Height)
{
	int FoundNodeIndex = -1;
	// Search through nodes to find the element to remove
	//@todo - traverse the tree instead of iterating through all nodes
	for (int NodeIndex = 0; NodeIndex < (int)m_Nodes.Size(); NodeIndex++)
	{
		AtlasNode&	Node = m_Nodes[NodeIndex];
		if (Node.OffsetX == OffsetX && Node.OffsetY == OffsetY && Node.Width == Width && Node.Height == Height)
		{
			FoundNodeIndex = NodeIndex;
			break;
		}
	}

	if (FoundNodeIndex != -1)
	{
		// Mark the found node as not being used anymore
		m_Nodes[FoundNodeIndex].bUsed = FALSE;

		// Walk up the tree to find the node closest to the root that doesn't have any used children
		INT ParentNodeIndex = FindParentNode(FoundNodeIndex);
		ParentNodeIndex = IsNodeUsed(ParentNodeIndex) ? INDEX_NONE : ParentNodeIndex;
		INT LastParentNodeIndex = ParentNodeIndex;
		while (ParentNodeIndex != INDEX_NONE 
			&& !IsNodeUsed(Nodes(ParentNodeIndex).ChildA) 
			&& !IsNodeUsed(Nodes(ParentNodeIndex).ChildB))
		{
			LastParentNodeIndex = ParentNodeIndex;
			ParentNodeIndex = FindParentNode(ParentNodeIndex);
		} 

		// Remove the children of the node closest to the root with only unused children,
		// Which restores the tree to its state before this element was allocated,
		// And allows allocations as large as LastParentNode in the future.
		if (LastParentNodeIndex != INDEX_NONE)
		{
			RemoveChildren(LastParentNodeIndex);
		}
		return TRUE;
	}

	return FALSE;
}